package com.cloud.configure;

import com.cloud.message.AbstractMQPushConsumer;
import com.cloud.message.AbstractMQTransactionProducer;
import com.cloud.message.MessageExtConst;
import com.cloud.rocketmq.MQConsumer;
import com.cloud.rocketmq.MQTransactionProducer;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.acl.common.AclClientRPCHook;
import org.apache.rocketmq.acl.common.SessionCredentials;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeOrderlyContext;
import org.apache.rocketmq.client.consumer.rebalance.AllocateMessageQueueAveragely;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.producer.DefaultMQProducer;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.remoting.RPCHook;
import org.apache.rocketmq.remoting.protocol.heartbeat.MessageModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.Assert;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

/**
 * RocketMq启动配置
 * 参考：https://github.com/TonyStark888/rocketmq-spring-boot-starter.git
 *
 * @author yzj
 */
@Configuration
public class RocketMqConfig implements ApplicationContextAware {

    private static Logger logger = LoggerFactory.getLogger(RocketMqConfig.class);

    @Resource
    private RocketMqProperties mqProperties;

    private ConfigurableApplicationContext applicationContext;

    private Map<String, TransactionMQProducer> producerMap = new HashMap<>();

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = (ConfigurableApplicationContext) applicationContext;
    }

    /**
     * RocketMQ 消费者初始化
     *
     * @return
     */
    @PostConstruct
    public void mqConsumer() {
        Map<String, Object> consumerBeans = applicationContext.getBeansWithAnnotation(MQConsumer.class);
        if (ObjectUtils.isNotEmpty(consumerBeans)) {
            for (Map.Entry<String, Object> entry : consumerBeans.entrySet()) {
                loadMQConsumer(entry.getKey(), entry.getValue());
            }
            Map<String, Object> producerBeans = applicationContext.getBeansWithAnnotation(MQTransactionProducer.class);
            for (Map.Entry<String, Object> entry : producerBeans.entrySet()) {
                initProducer(entry.getKey(), entry.getValue());
            }
        }
    }

    /**
     * RocketMQ 生产者对象初始化
     *
     * @return
     */
    @Bean
    public DefaultMQProducer mqProducer() {
        // 普通消息的生产者实例
        DefaultMQProducer producer = null;
        if (StringUtils.isNotEmpty(mqProperties.getAccount()) || StringUtils.isNotEmpty(mqProperties.getPassword())) {
            producer = new DefaultMQProducer(mqProperties.getProducerGroup(), initAclRPCHook());
        } else {
            producer = new DefaultMQProducer(mqProperties.getProducerGroup());
        }
        producer.setNamesrvAddr(mqProperties.getHostUrl());
        try {
            producer.start();
            logger.info("RocketMQ生产者:{},启动成功", producer.getProducerGroup());
        } catch (MQClientException e) {
            logger.error("RocketMQ生产者,启动失败", e);
        }
        return producer;
    }

    public TransactionMQProducer getProducer(String key) {
        TransactionMQProducer producer = producerMap.get(key);
        if (null == producer) {
            throw new NullPointerException("这个交易生产者" + key + "还没有初始化...");
        }
        return producer;
    }

    private void loadMQConsumer(String beanName, Object bean) {
        try {
            MQConsumer mqConsumer = applicationContext.findAnnotationOnBean(beanName, MQConsumer.class);
            if (StringUtils.isEmpty(mqProperties.getHostUrl())) {
                throw new RuntimeException("必须定义RocketMQ服务器地址");
            }
            Assert.notNull(mqConsumer.consumerGroup(), "必须定义消费者的消费者组");
            Assert.notNull(mqConsumer.topic(), "必须定义消费者的主题");
            Assert.notNull(mqConsumer.tag(), "必须定义消费者标签");
            if (!AbstractMQPushConsumer.class.isAssignableFrom(bean.getClass())) {
                throw new RuntimeException(bean.getClass().getName() + "-consumer未实现Consumer抽象类");
            }
            // 配置Push_Consumer初始化参数
            DefaultMQPushConsumer consumer = null;
            if (StringUtils.isNotEmpty(mqProperties.getAccount()) || StringUtils.isNotEmpty(mqProperties.getPassword())) {
                consumer = new DefaultMQPushConsumer(mqConsumer.consumerGroup(), initAclRPCHook(), new AllocateMessageQueueAveragely());
            } else {
                consumer = new DefaultMQPushConsumer(mqConsumer.consumerGroup());
            }
            consumer.setNamesrvAddr(mqProperties.getHostUrl());
            consumer.setMessageModel(MessageModel.valueOf(mqConsumer.messageMode()));
            consumer.setConsumeThreadMin(mqProperties.getThreadMin());
            consumer.setConsumeThreadMax(mqProperties.getThreadMax());
            consumer.setPullBatchSize(mqProperties.getPullBatchSize());
            consumer.setMaxReconsumeTimes(mqProperties.getMaxReconsumeTimes());
            consumer.subscribe(mqConsumer.topic(), StringUtils.join(mqConsumer.tag(), "||"));
            AbstractMQPushConsumer abstractMQPushConsumer = (AbstractMQPushConsumer) bean;
            if (MessageExtConst.CONSUME_MODE_CONCURRENTLY.equals(mqConsumer.consumeMode())) {
                consumer.registerMessageListener((List<MessageExt> list, ConsumeConcurrentlyContext consumeConcurrentlyContext) ->
                        abstractMQPushConsumer.dealMessage(list, consumeConcurrentlyContext));
            } else if (MessageExtConst.CONSUME_MODE_ORDERLY.equals(mqConsumer.consumeMode())) {
                consumer.registerMessageListener((List<MessageExt> list, ConsumeOrderlyContext consumeOrderlyContext) ->
                        abstractMQPushConsumer.dealMessage(list, consumeOrderlyContext));
            } else {
                throw new RuntimeException("未知消费模式!仅支持CONCURRENTLY和ORDERLY");
            }
            abstractMQPushConsumer.setConsumer(consumer);
            consumer.start();
            logger.info("RocketMQ消费者:{}开始", consumer.getConsumerGroup());
        } catch (Exception e) {
            logger.error("RocketMQ消费者启动失败", e);
        }
    }

    /**
     * 发送消息时携带账号密码
     *
     * @return
     */
    private RPCHook initAclRPCHook() {
        return new AclClientRPCHook(new SessionCredentials(mqProperties.getAccount(), mqProperties.getPassword()));
    }

    private void initProducer(String beanName, Object bean) {
        try {
            MQTransactionProducer producerListener = applicationContext.findAnnotationOnBean(beanName, MQTransactionProducer.class);
            if (StringUtils.isEmpty(mqProperties.getHostUrl())) {
                throw new RuntimeException("必须定义名称服务器地址");
            }
            Assert.notNull(producerListener.producerGroup(), "必须定义生产者的生产者组");
            Assert.notNull(producerListener.topic(), "必须定义生产者的主题");
            Assert.notNull(producerListener.tag(), "必须定义生产者标签");
            if (!AbstractMQTransactionProducer.class.isAssignableFrom(bean.getClass())) {
                throw new RuntimeException(bean.getClass().getName() + "--Listener未实现AbstractMQTransactionProducer抽象类");
            }
            ExecutorService executorService = new ThreadPoolExecutor(5, 100,
                    100, TimeUnit.SECONDS, new ArrayBlockingQueue<Runnable>(2000), new ThreadFactory() {
                @Override
                public Thread newThread(Runnable r) {
                    Thread thread = new Thread(r);
                    thread.setName("rocketmq-transaction-message-thread-" + mqProperties.getProducerGroup());
                    return thread;
                }
            });
            // 事务消息的生产者实例，需要在启动前设置TransactionListener
            TransactionMQProducer transactionProducer = null;
            String tProducerGroup = "Transaction" + mqProperties.getProducerGroup();
            if (StringUtils.isNotEmpty(mqProperties.getAccount()) || StringUtils.isNotEmpty(mqProperties.getPassword())) {
                transactionProducer = new TransactionMQProducer(tProducerGroup, initAclRPCHook());
            } else {
                transactionProducer = new TransactionMQProducer(tProducerGroup);
            }
            transactionProducer.setNamesrvAddr(mqProperties.getHostUrl());
            transactionProducer.setExecutorService(executorService);
            transactionProducer.setTransactionListener((AbstractMQTransactionProducer) bean);
            transactionProducer.start();
            producerMap.put(producerListener.topic() + "_" + producerListener.tag(), transactionProducer);
            logger.info("RocketMQ交易生产者:{}开始", producerListener.producerGroup());
        } catch (Exception e) {
            logger.error("RocketMQ交易生产者启动失败", e);
        }
    }

}
